import torch
import torchvision
from torch import nn
from torch.utils.tensorboard import SummaryWriter

input = torch.tensor([[1, -0.5],
                      [-1, 3]])

input = torch.reshape(input, (-1, 1, 2, 2))
# print(input.shape)

"""
root='./data/data_non_linear'：指定数据集下载和存储的根目录路径。这里数据集会被下载到当前目录下的data/data_non_linear文件夹中。

train=False：表示加载的是测试集（而不是训练集）。CIFAR-10的6万张图像中，5万张是训练集，1万张是测试集。

download=True：如果指定路径下没有数据集，会自动从互联网下载。

transform=torchvision.transforms.ToTensor()：指定数据转换方式：

ToTensor()将PIL图像或numpy数组转换为PyTorch张量(tensor)

同时会自动将像素值从[0,255]范围缩放到[0.0,1.0]范围

还会将图像维度从(H,W,C)变为(C,H,W)格式（适合PyTorch处理）
"""
datasets = torchvision.datasets.CIFAR10(root='./data/data_non_linear', train=False, download=True,
                                        transform=torchvision.transforms.ToTensor())

train_loader = torch.utils.data.DataLoader(datasets, batch_size=64)


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.relu1 = nn.ReLU()
        self.sigmoid1 = nn.Sigmoid()

    def forward(self, x):
        y = self.sigmoid1(x)
        return y


net = Net()
# output = net(input)
# print(output)

writer = SummaryWriter("logs/log_non_linear")
step = 0
for data in train_loader:
    img, label = data
    writer.add_images("input", img, step)
    output = net(img)
    writer.add_images("output", output, step)
    step = step + 1

writer.close()
